#!/usr/bin/env python3
import socket
import threading
import ssl
import time
import json
import logging
import base64
import subprocess
import hashlib
import struct
from urllib.parse import urlparse
from datetime import datetime

class IDEAProxy:
    def __init__(self, host='0.0.0.0', port=8080, upstream_proxy=None):
        self.host = host
        self.port = port
        self.running = False
        self.requests_log = []
        self.websocket_connections = {}  # Track active WebSocket connections
        
        # Upstream proxy configuration
        self.upstream_proxy = upstream_proxy  # {'host': str, 'port': int, 'username': str, 'password': str}
        
        logging.basicConfig(
            level=logging.INFO,
            format='%(asctime)s - %(levelname)s - %(message)s',
            handlers=[
                logging.FileHandler('idea_proxy.log'),
                logging.StreamHandler()
            ]
        )
        self.logger = logging.getLogger(__name__)

    def _create_upstream_connection(self, target_host, target_port):
        """Create connection through upstream proxy if configured"""
        if not self.upstream_proxy:
            # Direct connection
            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            sock.settimeout(10)
            sock.connect((target_host, target_port))
            return sock
        
        # Connect through upstream proxy
        proxy_host = self.upstream_proxy['host']
        proxy_port = self.upstream_proxy['port']
        proxy_username = self.upstream_proxy.get('username')
        proxy_password = self.upstream_proxy.get('password')
        
        # Connect to upstream proxy
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        sock.settimeout(10)
        sock.connect((proxy_host, proxy_port))
        
        # Send CONNECT request to upstream proxy
        connect_request = f"CONNECT {target_host}:{target_port} HTTP/1.1\r\n"
        connect_request += f"Host: {target_host}:{target_port}\r\n"
        
        # Add proxy authentication if provided
        if proxy_username and proxy_password:
            credentials = f"{proxy_username}:{proxy_password}"
            encoded_credentials = base64.b64encode(credentials.encode()).decode()
            connect_request += f"Proxy-Authorization: Basic {encoded_credentials}\r\n"
        
        connect_request += "\r\n"
        
        sock.send(connect_request.encode())
        
        # Read response from upstream proxy
        response = sock.recv(4096).decode()
        if not response.startswith("HTTP/1.1 200"):
            sock.close()
            raise Exception(f"Upstream proxy connection failed: {response.split()[1] if len(response.split()) > 1 else 'Unknown error'}")
        
        self.logger.info(f"Connected through upstream proxy {proxy_host}:{proxy_port} to {target_host}:{target_port}")
        return sock

    def _create_upstream_http_connection(self, target_host, target_port):
        """Create HTTP connection through upstream proxy"""
        if not self.upstream_proxy:
            # Direct connection
            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            sock.settimeout(10)
            sock.connect((target_host, target_port))
            return sock, False  # Not using proxy
        
        # Connect to upstream proxy for HTTP
        proxy_host = self.upstream_proxy['host']
        proxy_port = self.upstream_proxy['port']
        
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        sock.settimeout(10)
        sock.connect((proxy_host, proxy_port))
        
        self.logger.info(f"Connected to upstream HTTP proxy {proxy_host}:{proxy_port}")
        return sock, True  # Using proxy

    def start(self):
        self.server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self.server_socket.bind((self.host, self.port))
        self.server_socket.listen(5)
        
        self.running = True
        self.logger.info("=" * 60)
        self.logger.info(f"🚀 IDEA 代理服务器已启动")
        self.logger.info(f"🌐 监听地址: {self.host}:{self.port}")
        
        # WSL环境特殊提示
        if self.host == '0.0.0.0':
            try:
                # 获取WSL IP地址
                result = subprocess.run(['hostname', '-I'], capture_output=True, text=True, timeout=5)
                wsl_ip = result.stdout.strip().split()[0] if result.stdout.strip() else "未知"
                
                self.logger.info("💡 WSL环境配置:")
                self.logger.info(f"   当前WSL IP: {wsl_ip}")
                self.logger.info(f"   Windows下IDEA代理配置为: {wsl_ip}:8080")
                self.logger.info("   如IP获取失败，可在Windows中运行: wsl hostname -I")
            except:
                self.logger.info("💡 WSL环境配置:")
                self.logger.info("   Windows下IDEA代理配置为: <WSL_IP>:8080")
                self.logger.info("   获取WSL IP: wsl hostname -I")
        
        if self.upstream_proxy:
            self.logger.info(f"🔗 上游代理: {self.upstream_proxy['host']}:{self.upstream_proxy['port']}")
        else:
            self.logger.info("🔗 直连模式 (无上游代理)")
        self.logger.info("⏳ 等待客户端连接...")
        self.logger.info("=" * 60)
        
        try:
            while self.running:
                client_socket, addr = self.server_socket.accept()
                self.logger.info(f"🔗 新连接来自: {addr[0]}:{addr[1]}")
                
                client_thread = threading.Thread(
                    target=self.handle_client,
                    args=(client_socket, addr)
                )
                client_thread.daemon = True
                client_thread.start()
        except KeyboardInterrupt:
            self.stop()

    def stop(self):
        self.running = False
        if hasattr(self, 'server_socket'):
            self.server_socket.close()
        self.logger.info("=" * 60)
        self.logger.info("🛑 代理服务器已停止")
        self.logger.info(f"📊 总共处理了 {len(self.requests_log)} 个请求")  
        self.logger.info("=" * 60)

    def handle_client(self, client_socket, addr):
        try:
            request_data = client_socket.recv(4096).decode('utf-8')
            if not request_data:
                return
            
            self.log_request(request_data, addr)
            
            # Parse the request
            lines = request_data.split('\n')
            first_line = lines[0]
            parts = first_line.split(' ')
            if len(parts) < 3:
                self.logger.error(f"Invalid request format: {first_line}")
                return
                
            method, url, protocol = parts[0], parts[1], parts[2]
            
            # Handle CONNECT method for HTTPS
            if method == 'CONNECT':
                self.handle_https_connect(client_socket, url)
            else:
                # Check if this is a WebSocket upgrade request
                if self.is_websocket_request(request_data):
                    self.handle_websocket_request(client_socket, request_data, url)
                else:
                    self.handle_http_request(client_socket, request_data, url)
                
        except Exception as e:
            self.logger.error(f"Error handling client {addr}: {e}")
        finally:
            client_socket.close()

    def handle_https_connect(self, client_socket, url):
        try:
            if ':' in url:
                host, port = url.split(':')
                port = int(port)
            else:
                host = url
                port = 443  # Default HTTPS port
            
            self.logger.info(f"🔒 HTTPS CONNECT 请求: {host}:{port}")
            
            # Connect to target server through upstream proxy if configured
            target_socket = self._create_upstream_connection(host, port)
            
            # Send 200 Connection established
            client_socket.send(b'HTTP/1.1 200 Connection established\r\n\r\n')
            self.logger.info(f"✅ HTTPS 隧道已建立到 {host}:{port}")
            
            # Start tunneling
            self.tunnel_data(client_socket, target_socket, host)
            
        except Exception as e:
            self.logger.error(f"❌ HTTPS CONNECT 错误: {e}")
            client_socket.send(b'HTTP/1.1 500 Connection failed\r\n\r\n')

    def handle_http_request(self, client_socket, request_data, url):
        try:
            parsed_url = urlparse(url)
            host = parsed_url.netloc or parsed_url.path.split('/')[0]
            port = 80
            
            if ':' in host:
                host, port = host.split(':')
                port = int(port)
            
            # Connect to target server through upstream proxy if configured
            target_socket, using_proxy = self._create_upstream_http_connection(host, port)
            
            # Modify request for upstream proxy if needed
            if using_proxy and self.upstream_proxy:
                # Add proxy authentication to HTTP request if configured
                request_lines = request_data.split('\n')
                
                if self.upstream_proxy.get('username') and self.upstream_proxy.get('password'):
                    credentials = f"{self.upstream_proxy['username']}:{self.upstream_proxy['password']}"
                    encoded_credentials = base64.b64encode(credentials.encode()).decode()
                    
                    # Find where to insert Proxy-Authorization header
                    insert_index = 1  # After the first line (method line)
                    for i, line in enumerate(request_lines[1:], 1):
                        if line.strip() == '':  # Empty line indicates end of headers
                            insert_index = i
                            break
                    
                    # Insert Proxy-Authorization header
                    proxy_auth_header = f"Proxy-Authorization: Basic {encoded_credentials}"
                    request_lines.insert(insert_index, proxy_auth_header)
                    request_data = '\n'.join(request_lines)
            
            # Forward request
            target_socket.send(request_data.encode('utf-8'))
            
            # Get response
            response = b''
            while True:
                try:
                    chunk = target_socket.recv(4096)
                    if not chunk:
                        break
                    response += chunk
                    client_socket.send(chunk)
                except socket.timeout:
                    break
            
            self.log_response(response, host)
            
            target_socket.close()
            
        except Exception as e:
            self.logger.error(f"HTTP request error: {e}")
            error_response = b'HTTP/1.1 502 Bad Gateway\r\n\r\nProxy Error'
            client_socket.send(error_response)

    def is_websocket_request(self, request_data):
        """Check if the request is a WebSocket upgrade request"""
        request_lower = request_data.lower()
        return (
            'upgrade: websocket' in request_lower and
            'connection:' in request_lower and 
            'upgrade' in request_lower and
            'sec-websocket-key:' in request_lower
        )
    
    def handle_websocket_request(self, client_socket, request_data, url):
        """Handle WebSocket upgrade request"""
        try:
            parsed_url = urlparse(url)
            host = parsed_url.netloc or parsed_url.path.split('/')[0]
            port = 80
            
            if ':' in host:
                host, port = host.split(':')
                port = int(port)
            elif parsed_url.scheme == 'wss':
                port = 443
            elif parsed_url.scheme == 'ws':
                port = 80
                
            self.logger.info(f"🔌 WebSocket 升级请求: {host}:{port}")
            self.logger.info(f"🔌 URL: {url}")
            
            # Connect to target server
            target_socket, using_proxy = self._create_upstream_http_connection(host, port)
            
            # Forward the WebSocket upgrade request
            target_socket.send(request_data.encode('utf-8'))
            
            # Get the upgrade response
            response = b''
            while True:
                try:
                    chunk = target_socket.recv(4096)
                    if not chunk:
                        break
                    response += chunk
                    # WebSocket upgrade response should end with \r\n\r\n
                    if b'\r\n\r\n' in response:
                        break
                except socket.timeout:
                    break
            
            # Forward the response to client
            client_socket.send(response)
            
            # Check if upgrade was successful
            response_str = response.decode('utf-8', errors='ignore')
            if '101 switching protocols' in response_str.lower():
                self.logger.info(f"✅ WebSocket 连接已建立: {host}:{port}")
                
                # Track this WebSocket connection
                conn_id = f"{host}:{port}_{int(time.time())}"
                self.websocket_connections[conn_id] = {
                    'host': host,
                    'port': port,
                    'start_time': datetime.now().isoformat(),
                    'client_socket': client_socket,
                    'target_socket': target_socket
                }
                
                # Start WebSocket data tunneling with message monitoring
                self.tunnel_websocket_data(client_socket, target_socket, host, conn_id)
            else:
                response_first_line = response_str.split('\n')[0]
                self.logger.error(f"❌ WebSocket 升级失败: {response_first_line}")
                target_socket.close()
                
        except Exception as e:
            self.logger.error(f"❌ WebSocket 请求处理错误: {e}")
            error_response = b'HTTP/1.1 502 Bad Gateway\r\n\r\nWebSocket proxy error'
            client_socket.send(error_response)

    def tunnel_websocket_data(self, client_socket, target_socket, host, conn_id):
        """Tunnel WebSocket data with message monitoring"""
        def forward_websocket_data(source, destination, direction):
            try:
                total_bytes = 0
                message_count = 0
                while conn_id in self.websocket_connections:
                    try:
                        data = source.recv(4096)
                        if not data:
                            break
                        destination.send(data)
                        total_bytes += len(data)
                        
                        # Try to parse WebSocket frames for monitoring
                        try:
                            messages = self.parse_websocket_frames(data)
                            for msg in messages:
                                message_count += 1
                                self.log_websocket_message(host, direction, msg, message_count)
                        except:
                            # If parsing fails, just log raw data transfer
                            pass
                        
                        # Log periodic stats
                        if total_bytes > 0 and total_bytes % (1024*100) == 0:  # Every 100KB
                            self.logger.info(f"📊 WebSocket {host} {direction}: {total_bytes} 字节, {message_count} 消息")
                            
                    except socket.timeout:
                        continue
                    except Exception as e:
                        self.logger.info(f"🔚 WebSocket {host} {direction} 连接中断: {e}")
                        break
                        
                if total_bytes > 0:
                    self.logger.info(f"🔚 WebSocket {host} {direction} 连接结束: {total_bytes} 字节, {message_count} 消息")
                    
            except Exception as e:
                self.logger.info(f"🔚 WebSocket {host} {direction} 转发错误: {e}")
            finally:
                # Clean up connection tracking
                if conn_id in self.websocket_connections:
                    del self.websocket_connections[conn_id]

        # Start forwarding in both directions
        client_to_server = threading.Thread(
            target=forward_websocket_data,
            args=(client_socket, target_socket, "C->S")
        )
        server_to_client = threading.Thread(
            target=forward_websocket_data,
            args=(target_socket, client_socket, "S->C")
        )
        
        client_to_server.daemon = True
        server_to_client.daemon = True
        
        client_to_server.start()
        server_to_client.start()
        
        client_to_server.join()
        server_to_client.join()
        
        # Final cleanup
        if conn_id in self.websocket_connections:
            del self.websocket_connections[conn_id]
            
    def parse_websocket_frames(self, data):
        """Parse WebSocket frames to extract messages"""
        messages = []
        offset = 0
        
        while offset < len(data):
            if len(data) - offset < 2:
                break
                
            # Parse WebSocket frame header
            first_byte = data[offset]
            second_byte = data[offset + 1]
            
            fin = (first_byte & 0x80) != 0
            opcode = first_byte & 0x0F
            masked = (second_byte & 0x80) != 0
            payload_length = second_byte & 0x7F
            
            header_length = 2
            
            # Extended payload length
            if payload_length == 126:
                if len(data) - offset < 4:
                    break
                payload_length = struct.unpack('>H', data[offset + 2:offset + 4])[0]
                header_length = 4
            elif payload_length == 127:
                if len(data) - offset < 10:
                    break
                payload_length = struct.unpack('>Q', data[offset + 2:offset + 10])[0]
                header_length = 10
            
            # Masking key
            if masked:
                if len(data) - offset < header_length + 4:
                    break
                mask = data[offset + header_length:offset + header_length + 4]
                header_length += 4
            
            # Check if we have the complete frame
            if len(data) - offset < header_length + payload_length:
                break
                
            # Extract payload
            payload = data[offset + header_length:offset + header_length + payload_length]
            
            # Unmask if needed
            if masked:
                payload = bytes(payload[i] ^ mask[i % 4] for i in range(len(payload)))
            
            # Create message info
            message_info = {
                'opcode': opcode,
                'fin': fin,
                'length': payload_length,
                'payload': payload[:100] if len(payload) > 100 else payload  # Limit payload for logging
            }
            messages.append(message_info)
            
            offset += header_length + payload_length
            
        return messages
    
    def log_websocket_message(self, host, direction, message_info, message_count):
        """Log WebSocket message details"""
        opcode_names = {
            0x0: 'CONTINUATION',
            0x1: 'TEXT',
            0x2: 'BINARY', 
            0x8: 'CLOSE',
            0x9: 'PING',
            0xA: 'PONG'
        }
        
        opcode_name = opcode_names.get(message_info['opcode'], f"UNKNOWN({message_info['opcode']})")
        
        self.logger.info(f"📨 WebSocket {host} {direction} #{message_count}: {opcode_name} ({message_info['length']} bytes)")
        
        # Log text messages content (truncated)
        if message_info['opcode'] == 0x1:  # TEXT frame
            try:
                text_content = message_info['payload'].decode('utf-8')
                if len(text_content) > 200:
                    text_content = text_content[:200] + '...'
                self.logger.info(f"📨 内容: {text_content}")
            except:
                self.logger.info(f"📨 文本内容解码失败")

    def tunnel_data(self, client_socket, target_socket, host):
        def forward_data(source, destination, direction):
            try:
                total_bytes = 0
                while True:
                    data = source.recv(4096)
                    if not data:
                        break
                    destination.send(data)
                    total_bytes += len(data)
                    
                    # Log HTTPS traffic (encrypted) every 1MB or at connection end
                    if total_bytes % (1024*1024) == 0:
                        self.logger.info(f"📊 HTTPS {host} {direction}: {total_bytes} 字节传输")
                        
                if total_bytes > 0:
                    self.logger.info(f"🔚 HTTPS {host} {direction} 连接结束: 总计 {total_bytes} 字节")
            except Exception as e:
                self.logger.info(f"🔚 HTTPS {host} {direction} 连接中断: {e}")
                pass

        # Start forwarding in both directions
        client_to_server = threading.Thread(
            target=forward_data,
            args=(client_socket, target_socket, "C->S")
        )
        server_to_client = threading.Thread(
            target=forward_data,
            args=(target_socket, client_socket, "S->C")
        )
        
        client_to_server.daemon = True
        server_to_client.daemon = True
        
        client_to_server.start()
        server_to_client.start()
        
        client_to_server.join()
        server_to_client.join()

    def log_request(self, request_data, addr):
        timestamp = datetime.now().isoformat()
        
        lines = request_data.split('\n')
        if lines:
            method_line = lines[0].strip()
            headers = {}
            body = ''
            
            # 解析请求头
            header_end_index = 0
            for i, line in enumerate(lines[1:], 1):
                if line.strip() == '':  # 空行表示头部结束
                    header_end_index = i
                    break
                if ':' in line:
                    key, value = line.split(':', 1)
                    headers[key.strip()] = value.strip()
            
            # 提取请求体（如果有）
            if header_end_index > 0 and header_end_index + 1 < len(lines):
                body_lines = lines[header_end_index + 1:]
                body = '\n'.join(body_lines).strip()
            
            request_info = {
                'timestamp': timestamp,
                'client_addr': f"{addr[0]}:{addr[1]}",
                'method_line': method_line,
                'headers': headers,
                'body': body,
                'raw_request': request_data  # 保存完整的原始请求
            }
            
            self.requests_log.append(request_info)
            
            # Enhanced logging with more details
            self.logger.info("=" * 60)
            self.logger.info(f"新请求 - 时间: {timestamp}")
            self.logger.info(f"客户端地址: {addr[0]}:{addr[1]}")
            self.logger.info(f"请求方法行: {method_line}")
            
            # Log important headers
            important_headers = ['Host', 'User-Agent', 'Referer', 'Content-Type', 'Content-Length']
            self.logger.info("重要请求头:")
            for header in important_headers:
                if header in headers:
                    self.logger.info(f"  {header}: {headers[header]}")
            
            # Check if this looks like IDEA traffic
            user_agent = headers.get('User-Agent', '')
            if any(keyword in user_agent.lower() for keyword in ['idea', 'jetbrains', 'intellij']):
                self.logger.warning(f"⚠️  检测到IDEA流量: {method_line}")
                self.logger.warning(f"⚠️  User-Agent: {user_agent}")
            
            # Check if this is a WebSocket request
            if self.is_websocket_request(request_data):
                self.logger.warning(f"🔌 检测到WebSocket升级请求: {method_line}")
                self.logger.warning(f"🔌 Sec-WebSocket-Key: {headers.get('Sec-WebSocket-Key', 'N/A')}")
                self.logger.warning(f"🔌 Sec-WebSocket-Version: {headers.get('Sec-WebSocket-Version', 'N/A')}")
                if 'Sec-WebSocket-Protocol' in headers:
                    self.logger.warning(f"🔌 协议: {headers['Sec-WebSocket-Protocol']}")
            
            # Log all headers in debug mode
            if len(headers) > 0:
                self.logger.info("所有请求头:")
                for key, value in headers.items():
                    if key not in important_headers:
                        self.logger.info(f"  {key}: {value}")
            
            self.logger.info("=" * 60)

    def log_response(self, response_data, host):
        try:
            response_str = response_data.decode('utf-8', errors='ignore')
            
            # Extract status code from response
            lines = response_str.split('\n')
            status_line = lines[0] if lines else ''
            
            self.logger.info(f"📥 HTTP 响应来自 {host}: {len(response_data)} 字节")
            if status_line:
                self.logger.info(f"📥 状态: {status_line.strip()}")
                
            # Log response headers for debugging
            content_type = ''
            for line in lines[1:]:
                if line.lower().startswith('content-type:'):
                    content_type = line.split(':', 1)[1].strip()
                    self.logger.info(f"📥 内容类型: {content_type}")
                    break
                    
        except:
            self.logger.info(f"📥 HTTP 响应来自 {host}: {len(response_data)} 字节 (二进制数据)")

    def get_requests_summary(self):
        return {
            'total_requests': len(self.requests_log),
            'recent_requests': self.requests_log[-10:] if self.requests_log else [],
            'active_websockets': len(self.websocket_connections),
            'websocket_connections': list(self.websocket_connections.keys())
        }

    def save_logs_to_file(self, filename='idea_requests.json'):
        with open(filename, 'w') as f:
            json.dump(self.requests_log, f, indent=2)
        self.logger.info(f"Logs saved to {filename}")


if __name__ == '__main__':
    proxy = IDEAProxy()
    try:
        proxy.start()
    except KeyboardInterrupt:
        print("\nShutting down proxy...")
        proxy.stop()
        proxy.save_logs_to_file()